import torch
import torch.nn as nn

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

class BasicConv(nn.Module):
    def __init__(
        self,
        in_planes,
        out_planes,
        kernel_size,
        stride=1,
        padding=0,
        dilation=1,
        groups=1,
        relu=True,
        bn=True,
        bias=False,
    ):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(
            in_planes,
            out_planes,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )
        self.bn = (
            nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True)
            if bn
            else None
        )
        self.relu = nn.ReLU() if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x


class ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat(
            (torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1
        )


class SpatialGate(nn.Module):
    def __init__(self):
        super(SpatialGate, self).__init__()
        kernel_size = 7
        self.compress = ChannelPool()
        self.spatial = BasicConv(
            2, 1, kernel_size, stride=1, padding=(kernel_size - 1) // 2, relu=False
        )

    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.spatial(x_compress)
        scale = torch.sigmoid_(x_out)
        return x * scale

class reweight(nn.Module):
    def __init__(self):
        self.mlp = Mlp(6,10,3)

    def forward(self,x1,x2,x3):
        b,c,h,w=x1.shape
        x1_f = x1.reshape(b, -1)
        x2_f = x2.reshape(b, -1)
        x3_f = x3.reshape(b, -1)
        x1_in = torch.cat((x1_f.max(1)[0].reshape(b, -1),
                           x1_f.mean(1).reshape(b, -1)),
                          dim=1)
        x2_in = torch.cat((x2_f.max(1)[0].reshape(b, -1),
                           x2_f.mean(1).reshape(b, -1)),
                          dim=1)
        x3_in = torch.cat((x3_f.max(1)[0].reshape(b, -1),
                           x3_f.mean(1).reshape(b, -1)),
                          dim=1)
        mlp_input=torch.cat((torch.cat((x1_in,x2_in),dim=1),x3_in),dim=1)
        weight = self.mlp(mlp_input)


        out = (weight[:, 0]).reshape(b, 1, 1, 1) * x1 + (weight[:, 1]).reshape(b, 1, 1, 1) * x2 + (
        weight[:, 2]).reshape(b, 1, 1, 1) * x3

        return out

class TripletAttention(nn.Module):
    def __init__(
        self,
        gate_channels,
        reduction_ratio=16,
        pool_types=["avg", "max"],
        no_spatial=False,
    ):
        super(TripletAttention, self).__init__()
        self.ChannelGateH = SpatialGate()
        self.ChannelGateW = SpatialGate()
        self.no_spatial = no_spatial
        if not no_spatial:
            self.SpatialGate = SpatialGate()
            self.re = reweight()



    def forward(self, x):
        x_perm1 = x.permute(0, 2, 1, 3).contiguous()
        x_out1 = self.ChannelGateH(x_perm1)
        x_out11 = x_out1.permute(0, 2, 1, 3).contiguous()
        x_perm2 = x.permute(0, 3, 2, 1).contiguous()
        x_out2 = self.ChannelGateW(x_perm2)
        x_out21 = x_out2.permute(0, 3, 2, 1).contiguous()
        if not self.no_spatial:
            x_out = self.SpatialGate(x)
            # x_out = (1 / 3) * (x_out + x_out11 + x_out21)
            x_out = self.re(x_out,x_out11,x_out21)
        else:
            x_out = (1 / 2) * (x_out11 + x_out21)
        return x_out
